import cv2
import numpy as np
import glob

# 找棋盘格角点
# 阈值
criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
#棋盘格模板规格
w = 11   #内角点个数，内角点是和其他格子连着的点
h = 8

# 世界坐标系中的棋盘格点,例如(0,0,0), (1,0,0), (2,0,0) ....,(8,5,0)，去掉Z坐标，记为二维矩阵
objp = np.zeros((w*h,3), np.float32)
objp[:,:2] = np.mgrid[0:w,0:h].T.reshape(-1,2)*20
# 储存棋盘格角点的世界坐标和图像坐标对
objpoints = [] # 在世界坐标系中的三维点
imgpoints = [] # 在图像平面的二维点

images = glob.glob('img/*.jpg')
print(images)
for fname in images:
    img = cv2.imread(fname)
    gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
    ret, corners = cv2.findChessboardCorners(gray, (w,h),None)
    if ret == True:
        cv2.cornerSubPix(gray,corners,(11,11),(-1,-1),criteria)
        objpoints.append(objp)
        imgpoints.append(corners)
        cv2.drawChessboardCorners(img, (w,h), corners, ret)
        cv2.imshow('findCorners',img)
        cv2.waitKey(100)
    else:
        print(fname)
cv2.destroyAllWindows()
#标定、去畸变
print(gray.shape[::-1])
ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(objpoints, imgpoints, gray.shape[::-1], None, None)
# mtx：内参数矩阵
# dist：畸变系数
# rvecs：旋转向量 （外参数）
# tvecs ：平移向量 （外参数）
print (("ret:"),ret)
print (("mtx:\n"),mtx)
print (("dist:\n"),dist)
print (("rvecs:\n"),rvecs)
print (("tvecs:\n"),tvecs)
# 去畸变

img2 = cv2.imread('img/62.jpg')
h,w = img2.shape[:2]
newcameramtx, roi=cv2.getOptimalNewCameraMatrix(mtx,dist,(w,h),1,(w,h))
dst = cv2.undistort(img2, mtx, dist, None, newcameramtx)
x,y,w,h = roi
dst = dst[y:y+h, x:x+w]
cv2.imwrite('calibresult.jpg',dst)

# 反投影误差
total_error = 0
for i in range(len(objpoints)):
    imgpoints2, _ = cv2.projectPoints(objpoints[i], rvecs[i], tvecs[i], mtx, dist)
    error = cv2.norm(imgpoints[i],imgpoints2, cv2.NORM_L2)/len(imgpoints2)
    print(error)
    total_error += error
print (("total error: "), total_error/len(objpoints))

'''
mtx:
 [[894.348826     0.         936.44500723]
 [  0.         896.42857334 558.08348431]
 [  0.           0.           1.        ]]
dist:
 [[-0.00830184  0.10080361  0.00322837 -0.00452095 -0.1669461 ]]
'''
'''
mtx:
 [[1.03821817e+03 0.00000000e+00 1.00618003e+03]
 [0.00000000e+00 1.03840613e+03 4.94445799e+02]
 [0.00000000e+00 0.00000000e+00 1.00000000e+00]]
dist:
 [[-0.03792993  0.1240955   0.00080589 -0.00014446 -0.10270177]]
'''